package com.example.demo.util;
import com.sun.jna.Platform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.text.DecimalFormat;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class GPUUtil {

    public static List<Map<Object,Object>> getGPU() throws IOException {
        Process process = null;
        Process process2 = null;
        try {
            if (Platform.isWindows()) {
                process = Runtime.getRuntime().exec("nvidia-smi.exe");
                process2 = Runtime.getRuntime().exec("nvidia-smi.exe -L");
            } else if (Platform.isLinux()) {
                String[] shell = {"/bin/bash", "-c", "nvidia-smi"};
                process = Runtime.getRuntime().exec(shell);
                String[] shell2 = {"/bin/bash", "-c", "nvidia-smi -L"};
                process2 = Runtime.getRuntime().exec(shell2);
            }

            process.getOutputStream().close();
            process2.getOutputStream().close();
        } catch (IOException e) {
            e.printStackTrace();
            throw new IOException("显卡不存在或获取显卡信息失败");
        }

        BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
        BufferedReader reader2 = new BufferedReader(new InputStreamReader(process2.getInputStream()));

        StringBuffer stringBuffer = new StringBuffer();
        StringBuffer stringBuffer2 = new StringBuffer();
        String line = "";
        String line2 = "";
        while (null != (line = reader.readLine())) {
            stringBuffer.append(line + "\n");
        }
        while (null != (line2 = reader2.readLine())) {
            stringBuffer2.append(line2 + "\n");
        }
        String gpus = stringBuffer.toString();
        String gpuname = stringBuffer2.toString();
        String[] split = gpus.split("\\|===============================\\+======================\\+======================\\|");
        String[] gpusInfo = split[1].split("                                                                               ");
        String[] gpuInfo = gpusInfo[0].split("\\+-------------------------------\\+----------------------\\+----------------------\\+");
        List<Map<Object,Object>> data_list = new ArrayList<>();
        for (int i = 0; i < gpuInfo.length - 1; i++) {
            Map<Object, Object> data_map = new HashMap<>();
            String[] nameAndInfo = gpuInfo[i].split("\n");
            String[] split1 = nameAndInfo[1].split("\\|")[1] // 0  TITAN V             Off
                    .split("\\s+");//去空格

            data_map.put("number", Integer.parseInt(split1[1]));
            StringBuffer name = new StringBuffer();
            for (int j = 0; j < split1.length - 1; j++) {
                if (j > 1 && j != split1.length) {
                    name.append(split1[j] + " ");
                }
            }
            String[] split2 = gpuname.split(": ");
            for (int i1 = 0; i1 < split2.length; i1++) {
                if(split2[i1].indexOf(name.toString().replaceAll("\\.\\.\\. ","")) != -1){
                    data_map.put("name", split2[i1].split(" \\(UU")[0]);
                }
            }
            String[] info = nameAndInfo[2].split("\\|")[2].split("\\s+");
            int useable = Integer.parseInt(info[3].split("MiB")[0]) - Integer.parseInt(info[1].split("MiB")[0]);
            Double usageRate = Integer.parseInt(info[1].split("MiB")[0]) * 100.00 / Integer.parseInt(info[3].split("MiB")[0]);
            data_map.put("usageRate", usageRate);
            String c = nameAndInfo[2].split("\\|")[1].split("C")[0];
            data_map.put("usedMemory", formatByte(Integer.parseInt(info[1].split("MiB")[0])));
            data_map.put("totalMemory", formatByte(Integer.parseInt(info[3].split("MiB")[0])));
            data_map.put("useableMemory", formatByte(useable));
            if(c.indexOf("0%") != -1){
                c = c.replaceAll("0%","").replaceAll(" ","");
            }else if(c.indexOf("N/A") != -1){
                c = c.replaceAll("N/A","").replaceAll(" ","");
            }
            data_map.put("temperature",c);
            data_list.add(data_map);
        }
        return data_list;
    }

    public static String formatByte(long byteNumber){
        //换算单位
        double FORMAT = 1024.0;
        double gbNumber = byteNumber/FORMAT;
        return new DecimalFormat("#.#G").format(gbNumber);
    }
}
